import os
import typing

from openpyxl import Workbook
from openpyxl.worksheet.worksheet import Worksheet

from sai import Pipeline, Item


class ExcelPipeline(Pipeline):
    '''Buildin ExcelPipeline
    check excel_pipeline_usage_login2'''
    def __init__(self, fp: str=None):
        self.fp = fp
        self.wb: Workbook = None
        self.ws: Worksheet = None

    def open_spider(self, spider):
        downloads_dir = './downloads'
        if not os.path.exists(downloads_dir):
            os.makedirs(downloads_dir)
        if self.fp is None:
            self.fp = f"{downloads_dir}/{spider.name}.xlsx"

        self.wb = Workbook()
        self.ws = self.wb.active

        self.write_head = True

    def process_item(self, item: typing.Union[Item, dict], spider):
        if(self.write_head):
            excel_head = list(item.keys())
            self.ws.append(excel_head)
            self.write_head = False
        self.ws.append(list(item.values()))

    def close_spider(self, spider):
        self.wb.save(self.fp)
        self.wb.close()
